// implement fork from user space

#include <inc/string.h>
#include <inc/lib.h>

// PTE_COW marks copy-on-write page table entries.
// It is one of the bits explicitly allocated to user processes (PTE_AVAIL).
#define PTE_COW		0x800

//
// Custom page fault handler - if faulting page is copy-on-write,
// map in our own private writable copy.
//
static void
pgfault(struct UTrapframe *utf)
{
	void *addr = (void *) utf->utf_fault_va;
	uint32_t err = utf->utf_err;
	int r;

	// Check that the faulting access was (1) a write, and (2) to a
	// copy-on-write page.  If not, panic.
	// Hint:
	//   Use the read-only page table mappings at vpt
	//   (see <inc/memlayout.h>).

	// LAB 4: Your code here.
    if(!(vpt[PPN(addr)] & PTE_COW))
        panic("not COW");
    if(!((err & FEC_WR) && vpt[PPN(addr)] & PTE_COW))
        panic("not write fault");
    
	// Allocate a new page, map it at a temporary location (PFTEMP),
	// copy the data from the old page to the new page, then move the new
	// page to the old page's address.
	// Hint:
	//   You should make three system calls.
	//   No need to explicitly delete the old page's mapping.

	// LAB 4: Your code here.
    //envid_t envid = sys_getenvid();
    if(sys_page_alloc(0, PFTEMP, PTE_P|PTE_W|PTE_U) < 0)
    	panic("sys_page_alloc\n");
    memmove((void*)PFTEMP, (void *)ROUNDDOWN((unsigned int)addr, PGSIZE), PGSIZE);
    if(sys_page_map(0, (void *)PFTEMP, 0, (void *)ROUNDDOWN((unsigned int)addr, PGSIZE), PTE_W|PTE_U|PTE_P) != 0)
    	panic("sys_page_map\n");
    if(sys_page_unmap(0, PFTEMP) != 0)
    	panic("sys_page_unmap\n");
    return;
    
	//panic("pgfault not implemented");
}

//
// Map our virtual page pn (address pn*PGSIZE) into the target envid
// at the same virtual address.  If the page is writable or copy-on-write,
// the new mapping must be created copy-on-write, and then our mapping must be
// marked copy-on-write as well.  (Exercise: Why do we need to mark ours
// copy-on-write again if it was already copy-on-write at the beginning of
// this function?)
//
// Returns: 0 on success, < 0 on error.
// It is also OK to panic on error.
// 
static int
duppage(envid_t envid, unsigned pn)
{
	int r;
    pte_t pte = vpt[pn];
    void * add = (void *)(pn * PGSIZE);
    if(!(pte & PTE_P && pte & PTE_U))
        panic("perm invalid.");
    if(pte & PTE_W || pte & PTE_COW)
    {
        
        if ((r = sys_page_map(0, add, envid, add, PTE_P|PTE_U|PTE_COW)) < 0)
		return r;
        return sys_page_map(0, add, 0, add, PTE_P|PTE_U|PTE_COW);
    }
    
    return sys_page_map(0, add, envid, add, PTE_P|PTE_U);
    
    
	// LAB 4: Your code here.
	//panic("duppage not implemented");
	return 0;
}

//
// User-level fork with copy-on-write.
// Set up our page fault handler appropriately.
// Create a child.
// Copy our address space and page fault handler setup to the child.
// Then mark the child as runnable and return.
//
// Returns: child's envid to the parent, 0 to the child, < 0 on error.
// It is also OK to panic on error.
//
// Hint:
//   Use vpd, vpt, and duppage.
//   Remember to fix "env" in the child process.
//   Neither user exception stack should ever be marked copy-on-write,
//   so you must allocate a new page for the child's user exception stack.
//
envid_t
fork(void)
{
	// LAB 4: Your code here.
	//panic("fork not implemented");
    envid_t envid;
    envid = sys_exofork();
    set_pgfault_handler(pgfault);
    if(envid < 0)
        return envid;
    if(envid == 0)
    {
        // cprintf("CHILD\n");
        env = &envs[ENVX(sys_getenvid())];
        return 0;
    }
    
    //father
    unsigned int i, j, k = 0;
    //cprintf("%d,%d\n",UTOP,UXSTACKTOP);
    for(i = 0; i < VPD(UTOP); i ++)
    {
        if(vpd[i] == 0)
        {
            k += NPTENTRIES;
            continue;
        }
        else
        {
            for(j = 0; j < NPTENTRIES; j++, k++)
            {
                if(vpt[k] == 0)
                    continue;
                if(k == VPN(UXSTACKTOP) - 1)
                    continue;
                if(duppage(envid, k) < 0)
                    panic("duppage failed\n");
            }
        }
    }
    
        
    

    if(sys_page_alloc(envid, (void*)(UXSTACKTOP-PGSIZE), PTE_P|PTE_W|PTE_U) < 0)
        panic("alloc exception stack failed");
    extern void _pgfault_upcall(void);
    sys_env_set_pgfault_upcall(envid, _pgfault_upcall);
    if(sys_env_set_status(envid, ENV_RUNNABLE) < 0)
        panic("set status error on fork.");
    return envid;
    
    
}

// Challenge!
int
sfork(void)
{
	panic("sfork not implemented");
	return -E_INVAL;
}
